# import the necessary packages

from Common import settings
import flask
import redis
import time
import json


# initialize our Flask application and Redis server
app = flask.Flask(__name__)
db = redis.StrictRedis(host=settings.REDIS_HOST,
                       port=settings.REDIS_PORT, db=settings.REDIS_DB, password=settings.PASSWORD)

dbCloud = redis.StrictRedis(host=settings.Cloud_REDIS_HOST,
                       port=settings.Cloud_REDIS_PORT, db=settings.Cloud_REDIS_DB, password=settings.Cloud_PASSWORD)
# # 不同模型服务器的运行Redis队列名称
# IMAGE_QUEUE = {0: "resnet_queue", 1: "alexnet_queue", 2: "vggnet_queue"}

# 等待结果时的休眠时间
# CLIENT_SLEEP = 0.25


@app.route("/")
def homepage():
    return "Welcome to Edge Cloud Collarborative DNN Layer-Partitioning API!"


@app.route("/predict", methods=["POST"])
def predict():
    # initialize the data dictionary that will be returned from the
    # view
    data = {}
    # ensure an image was properly uploaded to our endpoint
    if flask.request.method == "POST":
        if flask.request.json != None:
            # read the image in PIL format and prepare it for
            # classification
            d = {"id": flask.request.json["id"], "image": flask.request.json["data"], "index": flask.request.json["index"], "ect": int(time.time()*1000)}
            db.rpush(settings.IMAGE_QUEUE[int(flask.request.json["appId"])], json.dumps(d))
            # keep looping until our model server returns the output
            # predictions
            k = flask.request.json["id"]
            while True:
                # attempt to grab the output predictions
                output = db.get(k)
                # check to see if our model has classified the input
                # image
                if output is not None:
                    # add the output predictions to our data
                    # dictionary so we can return it to the client
                    output = output.decode("utf-8")
                    output = json.loads(output)
                    data["predictions"] = output["result"]
                    data["et"] = output["et"]
                    data["tt"] = output["tt"]
                    data["ct"] = output["ct"]
                    # delete the result from the database and break
                    # from the polling loop
                    db.delete(k)
                    break
                # sleep for a small amount to give the model a chance
                # to classify the input image
                # time.sleep(settings.CLIENT_SLEEP)
            # indicate that the request was a success
            # data["success"] = True
            data["id"] = flask.request.json["id"]
            data["st"] = flask.request.json["st"]
            data["wt"] = flask.request.json["wt"]

            # data["endtime"] = int(time.time()*1000)
            data["rt"] = int(data["et"]) - int(data["st"])
            data["index"] = flask.request.json["index"]
            data["appId"] = flask.request.json["appId"]
            db.rpush("result", json.dumps(data))
    # return the data dictionary as a JSON response
    return flask.jsonify({"success": True})

# direct put the image to cloud process queue
@app.route("/cloudpredict", methods=["POST"])
def cloudpredict():
    # initialize the data dictionary that will be returned from the
    # view
    data = {}
    # ensure an image was properly uploaded to our endpoint
    if flask.request.method == "POST":
        if flask.request.json != None:
            # read the image in PIL format and prepare it for
            # classification
            d = {"id": flask.request.json["id"], "index": flask.request.json["index"], 
            "dtype": settings.IMAGE_DTYPE[int(flask.request.json["appId"])],
            "height": settings.IMAGE_HEIGHT[int(flask.request.json["appId"])], "width": settings.IMAGE_WIDTH[int(flask.request.json["appId"])],
            "chan": settings.IMAGE_CHANS[int(flask.request.json["appId"])], "temp": flask.request.json["data"], "tt": int(time.time()*1000), "ect": 0}
            # 发送到云中心的Redis
            dbCloud.rpush(settings.CLOUD_QUEUE[int(flask.request.json["appId"])], json.dumps(d))
            # keep looping until our model server returns the output
            # predictions
            k = flask.request.json["id"]
            while True:
                # 结果仍然写回边缘服务器
                # attempt to grab the output predictions
                output = db.get(k)
                # check to see if our model has classified the input
                # image
                if output is not None:
                    # add the output predictions to our data
                    # dictionary so we can return it to the client
                    output = output.decode("utf-8")
                    output = json.loads(output)
                    data["predictions"] = output["result"]
                    data["et"] = output["et"]
                    data["tt"] = output["tt"]
                    data["ct"] = output["ct"]
                    # delete the result from the database and break
                    # from the polling loop
                    db.delete(k)
                    break
                # sleep for a small amount to give the model a chance
                # to classify the input image
                # time.sleep(settings.CLIENT_SLEEP)
            # indicate that the request was a success
            # data["success"] = True
            data["id"] = flask.request.json["id"]
            data["st"] = flask.request.json["st"]
            data["wt"] = flask.request.json["wt"]
            data["rt"] = int(data["et"]) - int(data["st"])
            data["index"] = flask.request.json["index"]
            data["appId"] = flask.request.json["appId"]
            db.rpush("result", json.dumps(data))
    # return the data dictionary as a JSON response
    return flask.jsonify({"success": True})


# for debugging purposes, it's helpful to start the Flask testing
# server (don't use this for production
if __name__ == "__main__":
    print("* Starting web service...")
    app.run(host='0.0.0.0', port=17017)
